#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2022/6/23 上午10:40
# @Author  : qiufengfeng
# @File    : fgm_loss.py
# @Description :

import tensorflow as tf

@tf.function
def loss_with_gradient_penalty(model,loss_fuc,epsilon=1):
    '''
    在loss中添加 embedding层的干扰
    :param model:
    :param loss_fuc:
    :param epsilon:
    :return:
    '''
    def loss_with_gradient_penalty_2(y_true, y_pred):
        loss = tf.reduce_mean(loss_fuc(y_true, y_pred))
        embeddings = model.embedding.get_tokens_embeddings_variables()
        gp = tf.reduce_sum(tf.gradients(loss, [embeddings])[0].values**2)
        return loss + 0.5 * epsilon * gp
    return loss_with_gradient_penalty_2

